#!/usr/bin/env python
import os
import sys
from pathlib import Path


def setup_django_env():
    """设置Django环境"""
    # 获取项目根目录
    base_dir = Path(__file__).resolve().parent.parent

    # 将项目根目录添加到Python路径
    if str(base_dir) not in sys.path:
        sys.path.insert(0, str(base_dir))

    # 设置Django设置模块
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'datacharts.settings')

    # 初始化Django
    import django
    django.setup()


def read_and_import_excel(file_path):
    """读取Excel文件并导入到数据库"""
    import pandas as pd
    from accounts.models import StudentSurvey
    from django.db import transaction

    try:
        # 确保文件路径是绝对路径
        abs_file_path = Path(file_path).resolve()
        if not abs_file_path.exists():
            print(f"错误: 未找到文件 {abs_file_path}")
            return

        # 读取Excel文件
        print(f"正在读取文件: {abs_file_path}")
        df = pd.read_excel(abs_file_path)
        print("成功读取Excel文件!")

        # 显示数据基本信息
        print("\n数据基本信息:")
        print(df.info())

        # 数据清理
        df = df[df['age'].notna()]  # 过滤掉age为空的数据

        # 开始导入数据
        success_count = 0
        error_count = 0

        # 使用事务进行批量处理
        with transaction.atomic():
            for index, row in df.iterrows():
                try:
                    # 创建StudentSurvey实例
                    survey = StudentSurvey(
                        gradyear=int(row['gradyear']),
                        gender=row['gender'],
                        age=float(row['age']),
                        friends=int(row['friends']),
                        basketball=int(row.get('basketball', 0)),
                        football=int(row.get('football', 0)),
                        soccer=int(row.get('soccer', 0)),
                        softball=int(row.get('softball', 0)),
                        volleyball=int(row.get('volleyball', 0)),
                        swimming=int(row.get('swimming', 0)),
                        cheerleading=int(row.get('cheerleading', 0)),
                        baseball=int(row.get('baseball', 0)),
                        tennis=int(row.get('tennis', 0)),
                        sports=int(row.get('sports', 0)),
                        dance=int(row.get('dance', 0)),
                        band=int(row.get('band', 0)),
                        marching=int(row.get('marching', 0)),
                        music=int(row.get('music', 0)),
                        rock=int(row.get('rock', 0)),
                        god=int(row.get('god', 0)),
                        church=int(row.get('church', 0)),
                        jesus=int(row.get('jesus', 0)),
                        bible=int(row.get('bible', 0)),
                        hair=int(row.get('hair', 0)),
                        dress=int(row.get('dress', 0)),
                        blonde=int(row.get('blonde', 0)),
                        mall=int(row.get('mall', 0)),
                        shopping=int(row.get('shopping', 0)),
                        clothes=int(row.get('clothes', 0)),
                        hollister=int(row.get('hollister', 0)),
                        abercrombie=int(row.get('abercrombie', 0)),
                        cute=int(row.get('cute', 0)),
                        sexy=int(row.get('sexy', 0)),
                        hot=int(row.get('hot', 0)),
                        kissed=int(row.get('kissed', 0)),
                        drunk=int(row.get('drunk', 0)),
                        drugs=int(row.get('drugs', 0)),
                        death=int(row.get('death', 0)),
                        die=int(row.get('die', 0))
                    )
                    survey.save()
                    success_count += 1
                    if success_count % 100 == 0:  # 每导入100条数据打印一次进度
                        print(f"已成功导入 {success_count} 条数据")

                except Exception as e:
                    error_count += 1
                    print(f"导入第 {index + 1} 行数据时出错: {str(e)}")
                    continue

        print(f"\n导入完成!")
        print(f"成功导入: {success_count} 条")
        print(f"失败: {error_count} 条")

    except Exception as e:
        print(f"处理文件时发生错误: {str(e)}")


def main():
    """主函数"""
    # 设置Django环境
    setup_django_env()

    # 获取项目根目录
    base_dir = Path(__file__).resolve().parent.parent

    # 设置Excel文件路径
    excel_file = base_dir / 'data.xlsx'

    # 执行导入
    read_and_import_excel(excel_file)


if __name__ == '__main__':
    main()
